import os
import certifi
from paddlex import create_pipeline
# from paddleocr import PPChatOCRv4Doc
from typing import List, Dict
import base64
# 设置SSL证书路径
os.environ['SSL_CERT_FILE'] = certifi.where()

class OCRService:
    """
    OCR服务层类，封装文档识别与问答功能
    """
    def __init__(self, 
                 chat_api_key: str, 
                 retriever_api_key: str, 
                 base_url: str = "https://qianfan.baidubce.com/v2"):
        """
        初始化方法
        
        Args:
            chat_api_key: 聊天模块API Key
            retriever_api_key: 检索模块API Key
            base_url: 千帆服务基础URL（默认值为百度千帆服务地址）
        """
        self.chat_bot_config = {
            "module_name": "chat_bot",
            "model_name": "ernie-3.5-8k",
            "base_url": base_url,
            "api_type": "openai",
            "api_key": chat_api_key  # 聊天模块API Key
        }

        self.retriever_config = {
            "module_name": "retriever",
            "model_name": "embedding-v1",
            "base_url": base_url,
            "api_type": "qianfan",
            "api_key": retriever_api_key  # 检索模块API Key
        }

        # 初始化OCR管道（延迟加载预测器）
        self.pipeline =  create_pipeline(pipeline="PP-ChatOCRv4-doc", initial_predictor=False)

    

    def visual_analysis(self, 
                        file_path: str, 
                        use_orientation: bool = False,
                        use_unwarping: bool = False) -> List[Dict]:
        """
        文档图像视觉分析
        
        Args:
            file_path: 文件路径
            use_orientation: 是否启用文档方向分类（默认False）
            use_unwarping: 是否启用文档矫正（默认False）
            
        Returns:
            视觉分析结果列表
        """
        try:
           
            
            # 执行视觉预测
            visual_predict_res = self.pipeline.visual_predict(
                input=file_path,
                use_doc_orientation_classify=use_orientation,
                use_doc_unwarping=use_unwarping,
                use_common_ocr=True,
                use_seal_recognition=True,
                use_table_recognition=True
            )

            # 提取视觉信息
            visual_info_list = [res["visual_info"] for res in visual_predict_res]
            return visual_info_list

        except Exception as e:
            raise RuntimeError(f"视觉分析失败: {str(e)}") from e

    def generate_vector_info(self, visual_info_list: List[Dict]) -> Dict:
        """
        生成向量信息
        
        Args:
            visual_info_list: 视觉分析结果列表
            
        Returns:
            向量信息字典
        """
        try:
            return self.pipeline.build_vector(
                visual_info_list,
                flag_save_bytes_vector=True,
                retriever_config=self.retriever_config
            )

        except Exception as e:
            raise RuntimeError(f"向量生成失败: {str(e)}") from e

    def question_answer(self, 
                        keys: List[str], 
                        visual_info: List[Dict], 
                        vector_info: Dict) -> Dict:
        """
        文档内容问答
        
        Args:
            keys: 需要提取的关键字列表
            visual_info: 视觉分析结果
            vector_info: 向量信息
            
        Returns:
            问答结果字典
        """
        try:
            chat_result =  self.pipeline.chat(
                key_list=keys,
                visual_info=visual_info,
                vector_info=vector_info,
                chat_bot_config=self.chat_bot_config,
                retriever_config=self.retriever_config
            )
            print(chat_result)
            return chat_result

        except Exception as e:
            raise RuntimeError(f"问答处理失败: {str(e)}") from e

    def process_document(self, 
                         file_path: str, 
                         query_keys: List[str],
                         use_orientation: bool = False,
                         use_unwarping: bool = False) -> Dict:
        """
        文档处理全流程方法
        
        Args:
            file_path: 文件路径
            query_keys: 需要提取的关键字列表
            use_orientation: 是否启用文档方向分类（默认False）
            use_unwarping: 是否启用文档矫正（默认False）
            
        Returns:
            完整处理结果
        """
        try:
            # 1. 视觉分析
            visual_info = self.visual_analysis(
                file_path, 
                use_orientation, 
                use_unwarping
            )
            
            # 2. 生成向量信息
            vector_info = self.generate_vector_info(visual_info)
            
            # 3. 执行问答
            return self.question_answer(query_keys, visual_info, vector_info)
            
        except Exception as e:
            raise RuntimeError(f"文档处理失败: {str(e)}") from e
        
